import cv2
import numpy as np
import matplotlib.pyplot as plt


# 读取图像
image = cv2.imread('image.jpg')

# 转换为灰度图像
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

# 高斯模糊
blurred = cv2.GaussianBlur(gray, (5, 5), 0)

# 边缘检测
edges = cv2.Canny(blurred, 30, 100)

# 距离变换
dist_transform = cv2.distanceTransform(blurred, cv2.DIST_L2, 5)
ret, sure_bg = cv2.threshold(dist_transform, 0.7 * dist_transform.max(), 255, 0)

# 转换为uint8类型
sure_bg = np.uint8(sure_bg)

# 找到确定的前景
ret, sure_fg = cv2.threshold(blurred, 0.7 * blurred.max(), 255, 0)
sure_fg = np.uint8(sure_fg)

# 找到未知区域
unknown = cv2.subtract(sure_bg, sure_fg)

# 连通区域标记
ret, markers = cv2.connectedComponents(sure_fg)

# 为所有标记加1，确保背景为1
markers = markers + 1

# 未知区域标记为0
markers[unknown == 255] = 0

# 应用分水岭算法
markers = cv2.watershed(image, markers)

# 将分割结果标记为绿色
image[markers == -1] = [0, 255, 0]

# 显示结果
plt.figure(figsize=(10, 5))

plt.subplot(1, 2, 1)
plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
plt.title('Watershed Segmentation')

plt.subplot(1, 2, 2)
plt.imshow(markers, cmap='jet')
plt.title('Markers')

plt.show()

